from numpy import *

# 加载数据
def loadDataSet():
    postingList = [['my', 'dog', 'has', 'flea', 'problems', 'help', 'please'],
                   ['maybe', 'not', 'take', 'him', 'to', 'dog', 'park', 'stupid'],
                   ['my', 'dalmation', 'is', 'so', 'cute', 'I', 'love', 'him'],
                   ['stop', 'posting', 'stupid', 'worthless', 'garbage'],
                   ['mr', 'licks', 'ate', 'my', 'steak', 'how', 'to', 'stop', 'him'],
                   ['quit', 'buying', 'worthless', 'dog', 'food', 'stupid']]
    classVec = [0, 1, 0, 1, 0, 1]
    return postingList, classVec

# 合并所有单词，利用set来去重，得到所有单词的唯一列表
def createVocabList(dataSet):
    vocabSet = set([])
    for document in dataSet:
        vocabSet = vocabSet | set(document)
    return list(vocabSet)


# 优化词集模型= 为 词袋模型+=，将单词列表变为数字向量列表
def bagOfWords2VecMN(vocabList, inputSet):
    returnVec = [0] * len(vocabList)  # 获得所有单词等长的0列表
    for word in inputSet:
        if word in vocabList:
            returnVec[vocabList.index(word)] += 1  # 对应单词位置加1
    return returnVec


# 返回的是0、1各自两个分类中每个单词数量除以该分类单词总量再取对数ln 以及0、1两类的比例
def trainNB0(trainMatrix, trainCategory):
    numTrainDocs = len(trainMatrix)  # 样本数
    numWords = len(trainMatrix[0])  # 特征数
    pAbusive = sum(trainCategory) / float(numTrainDocs)  # 1类所占比例
    p0Num = ones(numWords)
    p1Num = ones(numWords)  # 初始化所有单词为1
    p0Denom = 2.0
    p1Denom = 2.0  # 初始化总单词为2        后面解释为什么这四个不初始化为0
    for i in range(numTrainDocs):
        if trainCategory[i] == 1:  # 求1类
            p1Num += trainMatrix[i]
            p1Denom += sum(trainMatrix[i])
        else:
            p0Num += trainMatrix[i]  # 求0类
            p0Denom += sum(trainMatrix[i])
    p1Vect = log(p1Num / p1Denom)  # numpy数组 / float = 1中每个单词/1中总单词
    p0Vect = log(p0Num / p0Denom)  # 这里为什么还用ln来处理，后面说明
    return p0Vect, p1Vect, pAbusive


# P(X|C)判断各类别的概率大小（这里是0、1）
def classifyNB(vec2Classify, p0Vec, p1Vec, pClass1):
    p1 = sum(vec2Classify * p1Vec) + log(pClass1)  # 相乘后得到哪些单词存在，再求和，再+log(P(C))
    p0 = sum(vec2Classify * p0Vec) + log(1.0 - pClass1)  # 由于使用的是ln，这里其实都是对数相加
    if p1 > p0:
        return 1
    else:
        return 0


# 封装调用的函数
def testingNB():
    listOPosts, listClasses = loadDataSet()
    myVocabList = createVocabList(listOPosts)
    trainMat = []
    for postinDoc in listOPosts:
        trainMat.append(bagOfWords2VecMN(myVocabList, postinDoc))
    p0V, p1V, pAb = trainNB0(array(trainMat), array(listClasses))
    # 上面求出了0、1两个类中各单词所占该类的比例，以及0、1的比例

    # 下面是预测两条样本数据的类别
    testEntry = ['love', 'my', 'dalmation']
    thisDoc = array(bagOfWords2VecMN(myVocabList, testEntry))  # 先将测试数据转为numpy的词袋模型 [0 2 0 5 1 0 0 3 ...]
    print(testEntry, 'classified as: ', classifyNB(thisDoc, p0V, p1V, pAb))  # 传值判断

    testEntry = ['stupid', 'garbage']
    thisDoc = array(bagOfWords2VecMN(myVocabList, testEntry))
    print(testEntry, 'classified as: ', classifyNB(thisDoc, p0V, p1V, pAb))


if __name__ == "__main__":
    testingNB()